#NOTE: 代码来源于FreeFEM导出的矩阵
# 矩阵主要为COO格式

import re
import numpy as np

def transmatrix(matrixdir, savematrix):
    i = 0
    base = 1

    for line in open(matrixdir): 
        line = line.strip("\n")
        if(line.startswith("% Size")):
            Size = re.compile("% Size = (\d+) (\d+)").match(line)
            numrow = int(Size.group(1))
            numcol = int(Size.group(2))
        elif(line.startswith("% Nonzeros")):
            Nonzeros = re.compile("% Nonzeros = (\d+)").match(line)
            nnz      = int(Nonzeros.group(1))
            row      = np.zeros(nnz, dtype=int)
            col      = np.zeros(nnz, dtype=int)
            val      = np.zeros(nnz, dtype=float)
        elif(line[0].isdigit()):
            line = re.split(r"[ ]+", line)
            row[i] = int(line[0]) - base
            col[i] = int(line[1]) - base
            val[i] = float(line[2])
            i += 1

    Z = zip(row,col,val)
    Z = sorted(Z,key=lambda s: s[0],reverse=False) 
    row,col,val = zip(*Z)

    row_count = np.zeros(numrow + 1, dtype=int)

    for i in range(0, len(row)):
        j = row[i]
        row_count[j] = row_count[j] + 1

    row_ptr = np.zeros(numrow + 1, dtype=int)
    for i in range(0, numrow):
        row_ptr[i+1] = row_count[i] + row_ptr[i]

    with open(savematrix, "w+" ) as fo:
        fo.write(str(numrow)+"\n")
        for i in range(0, numrow + 1):
            fo.write(str(row_ptr[i])+"\n")
        for i in range(0, len(col)):
            fo.write(str(col[i])+"\n")
        for i in range(0, len(val)):
            fo.write(str(val[i])+"\n")
    return numrow

def transrhs(numrow, rhsdir, saverhsdir):
    i = 0
    rhs = np.zeros(numrow, dtype=float)
    for line in open(rhsdir): 
        line = line.strip("\n")
        if(line[0].isdigit()):
            rhs[i] = line
            i += 1
    with open(saverhsdir, "w+" ) as fo:
        fo.write(str(numrow)+"\n")
        for i in range(0, numrow):
            fo.write(str(rhs[i])+"\n")

if __name__=="__main__":
    matrixdir = '/home/matrix/dxy/K4166006/A.txt'
    rhsdir    = '/home/matrix/dxy/K4166006/b.txt'
    savematrix = '/home/matrix/dxy/K4166006/A.JXPAMG'
    saverhsdir = '/home/matrix/dxy/K4166006/b.JXPAMG'
    numrow = transmatrix(matrixdir, savematrix)
    transrhs(numrow, rhsdir, saverhsdir)